{-# LANGUAGE OverloadedStrings, RecordWildCards, MultiWayIf #-}
module Main where

import Data.Int (Int64, Int16)
import Control.Exception (Exception, throwIO)
import Control.Concurrent (forkIO)
import Control.Monad (forM_, replicateM_)
import Text.Printf (printf)
import qualified System.IO.Streams.TCP as TCP
import qualified System.IO.Streams as Streams
import qualified Data.Connection as Conn
import qualified Database.PostgreSQL.Protocol.Readers as Readers
import qualified Database.PostgreSQL.Protocol.Builders as Builders
import qualified Database.PostgreSQL.Protocol.Types as PG
import qualified Database.PostgreSQL.Protocol.ObjectIDs as OID
import qualified Data.ByteString.Char8 as B
import qualified Data.ByteString.Builder as B
import qualified Data.ByteString.Lazy as Lazy
import qualified Data.Array as Array

tableOID = 10240

defaultFields :: [PG.FieldDescription]
defaultFields =
    [ PG.FieldDescription
        { PG.fieldName = "id"
        , PG.fieldTableID = tableOID
        , PG.fieldColumnID = 0
        , PG.fieldDataTypeID = OID.INTEGER
        , PG.fieldDataTypeSize = 4
        , PG.fieldDataTypeModifier = 0
        , PG.fieldFormat = PG.BinaryFormat
        }
    , PG.FieldDescription
        { PG.fieldName = "username"
        , PG.fieldTableID = tableOID
        , PG.fieldColumnID = 1
        , PG.fieldDataTypeID = OID.TEXT
        , PG.fieldDataTypeSize = -1
        , PG.fieldDataTypeModifier = 0
        , PG.fieldFormat = PG.TextFormat
        }
    , PG.FieldDescription
        { PG.fieldName = "password"
        , PG.fieldTableID = tableOID
        , PG.fieldColumnID = 2
        , PG.fieldDataTypeID = OID.TEXT
        , PG.fieldDataTypeSize = -1
        , PG.fieldDataTypeModifier = 0
        , PG.fieldFormat = PG.TextFormat
        }

    ]

defaultValues :: [[PG.Value]]
defaultValues =
    [ [ Just (B.toLazyByteString (B.int32BE 1))
      , Just "test user 1"
      , Just "password"
      ]
    , [ Just (B.toLazyByteString (B.int32BE 2))
      , Just "test user 2"
      , Just "password"
      ]
    ]

listToArray :: [a] -> PG.Array16 a
listToArray xs =
    let len = fromIntegral $ length xs
     in Array.array (0, len-1) (zip [0..len-1] xs)

data ConnectionClosedException = ConnectionClosedException
    deriving Show
instance Exception ConnectionClosedException

data NotImplementException = NotImplementException String
    deriving Show
instance Exception NotImplementException

readExact :: Streams.InputStream B.ByteString -> Int64 -> IO Lazy.ByteString
readExact stream size = do
    s <- readEnough stream (fromIntegral size)
    if Lazy.length s > size
      then do
          let (s', remaining) = Lazy.splitAt size s
          Streams.unRead (Lazy.toStrict remaining) stream
          return s'
      else return s

readEnough :: Streams.InputStream B.ByteString -> Int -> IO Lazy.ByteString
readEnough stream min_size = go 0 []
  where
      go :: Int -> [B.ByteString] -> IO Lazy.ByteString
      go size xs
        | size >= min_size = return (Lazy.fromChunks xs)
        | otherwise = do
            bs <- Streams.read stream >>= maybe (throwIO ConnectionClosedException) return
            go (size + B.length bs) (bs : xs)

sendBuilder :: TCP.TCPConnection -> B.Builder -> IO ()
sendBuilder conn builder = Conn.send conn $ B.toLazyByteString builder

simpleError :: Lazy.ByteString -> Lazy.ByteString -> PG.NoticeFields
simpleError code err = [ (PG.NoticeSeverity, "ERROR")
                  , (PG.NoticeCode, code)
                  , (PG.NoticeMessage, err)
                  ]

sendError :: TCP.TCPConnection -> PG.NoticeFields -> IO ()
sendError conn errors =
    sendBuilder conn $ Builders.errorResponseMessage errors

data ResultSet = ResultSet
    { resultFields :: [PG.FieldDescription]
    , resultRows :: [[PG.Value]]
    }

data CommandResult = CommandError PG.NoticeFields | CommandComplete (Maybe ResultSet) Lazy.ByteString

sendCommandResult conn (CommandComplete opt_result msg) = do
    case opt_result of
        Just ResultSet{..} -> do
            sendBuilder conn $ Builders.rowDescriptionMessage (listToArray resultFields)
            forM_ resultRows $ \row ->
                sendBuilder conn $ Builders.dataRowMessage $
                    listToArray row
        _ -> return ()
    sendBuilder conn $ Builders.commandCompleteMessage msg
sendCommandResult conn (CommandError errors) = sendError conn errors

sendCommandResultParse conn (CommandComplete results msg) = return ()
sendCommandResultParse conn (CommandError errors) = sendError conn errors

sendCommandResultDescribe conn (Just (CommandComplete (Just ResultSet{..}) _)) =
    sendBuilder conn $ Builders.rowDescriptionMessage (listToArray resultFields)
sendCommandResultDescribe conn _ = sendError conn $ simpleError "42P14" "invalid query"

sendCommandResultExecute conn (Just (CommandComplete (Just ResultSet{..}) _)) =
    forM_ resultRows $ \row ->
                sendBuilder conn $ Builders.dataRowMessage $
                    listToArray row
sendCommandResultExecute conn _ = sendError conn $ simpleError "42P14" "invalid query"

type QueryHandler = Lazy.ByteString -> IO CommandResult

defaultHandler sql = do
    if | "insert" `Lazy.isPrefixOf` sql ->
         return $ CommandComplete Nothing "INSERT 0 1"
       | "update" `Lazy.isPrefixOf` sql ->
         return $ CommandComplete Nothing "UPDATE 1"
       | "error" `Lazy.isPrefixOf` sql ->
         return $ CommandError (simpleError "23503" "foreign key violation")
       | otherwise ->
         return $ CommandComplete
             (Just ResultSet
             { resultFields = defaultFields
             , resultRows = defaultValues
             })
             "SELECT 2"

handle :: QueryHandler -> TCP.TCPConnection -> IO ()
handle handler conn = do
    startup
    loop True Nothing
  where
    stream = Conn.source conn
    startup = do
      msg <- Readers.readSessionMessage (readExact stream . fromIntegral)
      case msg of
          PG.SSLRequest -> sendBuilder conn Builders.sslRequestRejectedMessage >> startup
          PG.StartupMessage m n args -> do
              printf "%d.%d %s\n" m n (show args)
              sendBuilder conn $ Builders.authenticationResponseMessage PG.AuthenticationOk
              sendBuilder conn $ Builders.parameterStatusMessage "server_version" "11.0"
              sendBuilder conn $ Builders.parameterStatusMessage "client_encoding" "utf-8"
              sendBuilder conn $ Builders.readyForQueryMessage PG.TransactionIdle
          _ -> throwIO $ NotImplementException (show msg)
    loop False _ = return ()
    loop True result = do
      msg <- Readers.readFrontendMessage (readExact stream . fromIntegral)
      (continue, result') <- case msg of
          PG.Query sql -> do
              print $ "query: " <> sql
              result <- handler sql
              sendCommandResult conn result
              sendBuilder conn $ Builders.readyForQueryMessage PG.TransactionIdle
              return (True, Nothing)
          PG.Parse name sql params -> do
              print $ "parse: " <> name <> " " <> Lazy.toStrict sql
              result <- handler sql
              sendCommandResultParse conn result
              return (True, Just result)
          PG.Bind name statname params values fmt -> do
              print $ "bind: " <> statname
              sendBuilder conn $ Builders.bindCompleteMessage
              return (True, result)
          PG.Describe PG.StatementObject name -> do
              print $ "describe: " <> name
              sendCommandResultDescribe conn result
              return (True, result)
          PG.Describe PG.PortalObject name -> do
              print $ "describe: " <> name
              sendCommandResultDescribe conn result
              return (True, result)
          PG.Execute name n -> do
              print $ "execute: " <> name <> ", limit: " <> B.pack (show n)
              sendCommandResultExecute conn result
              return (True, Nothing)
          PG.Sync -> do
              sendBuilder conn $ Builders.readyForQueryMessage PG.TransactionIdle
              return (True, result)
          PG.Terminate -> return (False, Nothing)
          _ -> throwIO $ NotImplementException (show msg)
      loop continue result'


main :: IO ()
main = do
    sock <- TCP.bindAndListen 1024 3000
    go sock
  where
      go sock = do
          conn <- TCP.accept sock
          forkIO (handle defaultHandler conn)
          go sock
